#!/usr/bin/env bash

# Generate monthly accounting statistics from Slurm.
# Specific accounts and start/end dates may be specified.
# Homepage: https://github.com/OleHolmNielsen/Slurm_tools/

# CONFIGURE THIS for your local environment:
# People who will receive the report (default: Off)
export SEND_EMAIL=0
SUBSCRIBERS=root@localhost
# Report file prefix
REPORT_PREFIX=/tmp/Slurm_report
# Change the date/time format in report header for readability (formats in "man strftime")
export SLURM_TIME_FORMAT="%d-%b-%Y_%R"
# End of CONFIGURE lines

#
# Specify the accounting reports
#
FLAGS="-t hourper --tres=cpu,gpu"
TYPE="AccountUtilizationByUser"
OPTIONS="tree"

# Command usage:
function usage()
{
	cat << EOF

Usage: $0 [-m] [-a accountlist] [-s startdate -e enddate] [-r report-directory]
where:
	-m: Send report by E-mail to $SUBSCRIBERS
	-a: Select an account (or list of accounts)
	-s -e: Select Start and End dates of the report
	-c: Select current month from the 1st until today
	-w: Select the last week
	-r: Copy the report to a specified directory as well

Date format: MMDD

EOF
}

#####################################################################################
#
# Parse command options
#

while getopts "a:s:e:r:mcwh" options; do
	case $options in
		a )	export accountlist="Accounts=$OPTARG"
			echo Select accounts $OPTARG
			export REPORT_PREFIX=${REPORT_PREFIX}_${OPTARG}
			;;
		s )	export START="$OPTARG"
			echo Start date $OPTARG
			REPORT_NAME=${START}_${END}
			;;
		e )	export END="$OPTARG"
			echo End date $OPTARG
			REPORT_NAME=${START}_${END}
			;;
		c )	export START=`date +%m01%y`
			export END=`date +%m%d%y`
			echo Select current month from $START until $END
			REPORT_NAME=Current_month
			;;
		w )	export START=`date -d "last week" +%m%d%y`
			export END=`date +%m%d%y`
			echo Select last week from $START until $END
			REPORT_NAME=Last_week
			;;
		r )	export REPORTDIR="$OPTARG"
			echo Copy report to directory $OPTARG
			;;
		m )	export SEND_EMAIL=1
			echo Send report by E-mail to $SUBSCRIBERS
			;;
		V ) echo $my_version
			exit 1;;
		h|? ) usage
			exit 1;;
		* ) usage
			exit 1;;
	esac
done

# Test for extraneous command line arguments
if test $# -gt $(($OPTIND-1))
then
	echo ERROR: Too many command line arguments: $*
	usage
	exit 1
fi

#
# Default period: last month 
#
# Test if either START or END are empty strings
if test -z "$START" -o -z "$END"
then
	MONTH=`date -d "last month" +%B`
	YEAR=`date -d "last month" +%Y`
	REPORT_NAME=${MONTH}_${YEAR}
	# START=`date -d "-1 month -$(($(date +%-d)-1)) day" +%m%d%y`
	START=`date -d "last month" +%m01%y`
	# END=`date -d "-$(date +%d) day" +%D`
	END=`date +%m01%y`
fi

# Generate accounting report

export SYSTEMNAME=`scontrol show config | grep ClusterName`

REPORT=$REPORT_PREFIX.$REPORT_NAME

# First report overall statistics 
sreport cluster utilization Start=$START End=$END -t percent > $REPORT
# Report account usage
echo >> $REPORT
sreport $FLAGS cluster $TYPE $accountlist Start=$START End=$END $OPTIONS >> $REPORT

echo Report generated to file $REPORT

# Mail monthly report to a list of users (acct_list in /etc/aliases)
if test $SEND_EMAIL -gt 0
then
	cat ${REPORT} | /bin/mail -s "$SYSTEMNAME Slurm batch job accounting for $START until $END" ${SUBSCRIBERS}
fi

if test -n "$REPORTDIR"
then
	# Copy report to a report directory as well
	if test -d $REPORTDIR
	then
		echo Copy ${REPORT} to $REPORTDIR
		cp -p ${REPORT} $REPORTDIR
	else
		echo Report directory $REPORTDIR not accessible!
		ls -lad $REPORTDIR
	fi
fi
